package dbbench

import (
	"context"
	"database/sql"
	"encoding/base64"
	"github.com/VictoriaMetrics/fastcache"
	"github.com/go-redis/redis/v8"
	_ "github.com/go-sql-driver/mysql"
	"github.com/tikv/client-go/v2/config"
	"github.com/tikv/client-go/v2/rawkv"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type KVStore interface {
	Put(k, v []byte) error
	Get(k []byte) ([]byte, error)
	Delete(k []byte) error
}

type mysqlKvStore struct {
	db   *sql.DB
	stmt *sql.Stmt
}

func NewMysqlKvStore(host string) (*mysqlKvStore, error) {
	db, err := sql.Open("mysql", host)
	if err != nil {
		return nil, err
	}
	if err = db.Ping(); err != nil {
		return nil, err
	}
	stmt, err := db.Prepare(`insert into kv(id,val) values (?,?) ON DUPLICATE KEY UPDATE val=values(val)`)
	if err != nil {
		return nil, err
	}
	return &mysqlKvStore{
		db:   db,
		stmt: stmt,
	}, nil
}
func (m *mysqlKvStore) Put(k, v []byte) error {
	_, err := m.stmt.Exec(base64.StdEncoding.EncodeToString(k), base64.StdEncoding.EncodeToString(v))
	return err
}

func (m *mysqlKvStore) Get(k []byte) ([]byte, error) {
	//TODO implement me
	panic("implement me")
}

func (m *mysqlKvStore) Delete(k []byte) error {
	//TODO implement me
	panic("implement me")
}

type fastCacheKvStore struct {
	c *fastcache.Cache
}

func NewFastCacheKvStore() *fastCacheKvStore {
	c := fastcache.New(1024 * 1024 * 1024 * 2)
	return &fastCacheKvStore{c: c}
}

func (f *fastCacheKvStore) Put(k, v []byte) error {
	f.c.Set(k, v)
	return nil
}

func (f *fastCacheKvStore) Get(k []byte) ([]byte, error) {
	//TODO implement me
	panic("implement me")
}

func (f *fastCacheKvStore) Delete(k []byte) error {
	//TODO implement me
	panic("implement me")
}

type mongodbKVStore struct {
	//client *mongo.Client
	//db     *mongo.Database
	col *mongo.Collection
	raw bool
}

func NewMongodbKVStore(host, dbname string, raw bool) (*mongodbKVStore, error) {
	opt := options.Client().ApplyURI(host)
	client, err := mongo.NewClient(opt)
	if err != nil {
		return nil, err
	}
	if err = client.Connect(context.TODO()); err != nil {
		return nil, err
	}
	if err = client.Ping(context.Background(), nil); err != nil {
		return nil, err
	}
	col := client.Database(dbname).Collection("kv")
	return &mongodbKVStore{col: col, raw: raw}, err
}

func (m *mongodbKVStore) buildPk(k []byte) interface{} {
	return bson.M{"_id": string(k)}
}

type mongodbVal struct {
	Val []byte
}

func (m *mongodbKVStore) Put(k, v []byte) error {
	var val interface{}
	if m.raw {
		val = v
	} else {
		val = mongodbVal{Val: v}
	}
	_, err := m.col.ReplaceOne(context.TODO(), m.buildPk(k), val, options.Replace().SetUpsert(true))
	return err
}

func (m *mongodbKVStore) Get(k []byte) ([]byte, error) {
	findRet := m.col.FindOne(context.TODO(), m.buildPk(k))
	if err := findRet.Err(); err != nil {
		if err == mongo.ErrNoDocuments {
			return nil, nil
		}
		return nil, err
	}
	if m.raw {
		val, err := findRet.DecodeBytes()
		return val, err
	} else {
		v := mongodbVal{}
		err := findRet.Decode(&v)
		if err != nil {
			return nil, err
		}
		return v.Val, nil
	}

}

func (m *mongodbKVStore) Delete(k []byte) error {
	_, err := m.col.DeleteOne(context.TODO(), m.buildPk(k))
	return err
}

type tiKvStore struct {
	cli *rawkv.Client
}

func NewTiKvStore(host string) (*tiKvStore, error) {
	cli, err := rawkv.NewClient(context.TODO(), []string{host}, config.DefaultConfig().Security)
	if err != nil {
		return nil, err
	}
	return &tiKvStore{cli: cli}, nil
}

func (t *tiKvStore) Put(k, v []byte) error {
	return t.cli.Put(context.TODO(), k, v)
}

func (t *tiKvStore) Get(k []byte) ([]byte, error) {
	return t.cli.Get(context.TODO(), k)
}

func (t *tiKvStore) Delete(k []byte) error {
	return t.cli.Delete(context.TODO(), k)
}

type redisKVStore struct {
	cli *redis.Client
}

func NewRedisKVStore(host string) (*redisKVStore, error) {
	cli := redis.NewClient(&redis.Options{
		Addr: host,
	})
	return &redisKVStore{cli: cli}, nil
}

func (r *redisKVStore) Put(k, v []byte) error {
	cmd := r.cli.Set(context.TODO(), string(k), v, 0)
	return cmd.Err()
}

func (r *redisKVStore) Get(k []byte) ([]byte, error) {
	return nil, nil
}

func (r *redisKVStore) Delete(k []byte) error {
	return nil
}
